from llm_engine.config import CHROMA_DB_PATH
import chromadb
import numpy as np
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llm_engine.config import API_KEY, EMBEDDING_MODEL_NAME


def test_chromadb():
    '''
    空的
    '''
    client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
    print(len(client.list_collections()))
    print(client.list_collections()[0].name)
    collection = client.get_or_create_collection("project_collection")
    collection._embedding_dimension = 1024  # 运行到前一步的时候, 没有_embedding_dimension这个属性.
    print(collection.count())
    all_docs = collection.get(include=["embeddings", "documents", "metadatas", "uris"]) # 大模型可能让你读ids, 新版本中不让指定这个, 但是会自动返回ids.
    print(all_docs)
    
def test_chromadb2():
    '''
    手动写入能写, 说明智谱返回的向量不是1536维, 不符合openai规范, 被沉默抛弃了.
    
    手动调用add方法可以看到:
    <chromadb.utils.embedding_functions.DefaultEmbeddingFunction object at 0x0000026F414173E0>
    Traceback (most recent call last):
    File "d:\workspace\nisp-llm-poc\llm_engine\build_index.py", line 147, in <module>
        main()
    File "d:\workspace\nisp-llm-poc\llm_engine\build_index.py", line 98, in main
        chroma_collection.add(
    File "D:\workspace\nisp-llm-poc\env\Lib\site-packages\chromadb\api\models\Collection.py", line 89, in add
        self._client._add(
    File "D:\workspace\nisp-llm-poc\env\Lib\site-packages\chromadb\api\rust.py", line 407, in _add
        return self.bindings.add(
            ^^^^^^^^^^^^^^^^^^
    chromadb.errors.InvalidArgumentError: Collection expecting embedding with dimension of 1536, got 1024

    '''
    client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
    collection = client.get_or_create_collection("project_collection")
    # collection._embedding_dimension = 1024 # 没用, 一样Collection expecting embedding with dimension of 1536, got 1024
    embedding = np.random.rand(1024).tolist()  # 必须和 collection 的 dimension 一致
    collection.add(
        documents=["测试文档"],
        metadatas=[{"test": 1}],
        ids=["test_doc"],
        embeddings=[embedding]
    )
    
    all_docs = collection.get(include=["documents","metadatas","embeddings"])
    print(all_docs)
    
    
def test_embedding():
    '''
    确认智谱embedding模型返回的维度.
    '''
    embed_model = ZhipuAIEmbedding(
        api_key=API_KEY,
        model=EMBEDDING_MODEL_NAME
    )
    vector = embed_model.get_text_embedding("一根筋是导致两头堵的核心原因")
    print(type(vector), len(vector))  # 检查类型和长度  完犊子, 是1024. 

if __name__ == "__main__":
    test_chromadb()
    